using System;
using System.Collections.Generic;
using System.Linq;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics.Functions;
using Atomic.Libraries.Mathematics.Functions.DualNumbers;
using Atomic.Libraries.Mathematics.Optimization.Ipopt;
using Atomic.Libraries.Mathematics.RandomNumbers;
using Atomic.Structures;

namespace Atomic.Thermodynamics.PartitionFunctions
{
	public partial class IsobaricPartitionFunction
	{
		[Serializable]
		public class BivariateIsobaricPartitionFunction : IGibbsGrandPotential
		{
			private Atom atom1, atom2;
			private double[] counts1, counts2;
			private double count1_min, count1_max, count2_min, count2_max;
			private RollingCache<HashedTuple<double, double>, InnerSolver> solvers;

			public BivariateIsobaricPartitionFunction(IEnumerable<IIsobaricPartitionFunctionEntry> entries)
			{
				if (entries == null)
				{
					throw new ArgumentNullException();
				}

				Entries = new ReadOnlyList<IIsobaricPartitionFunctionEntry>(entries);

				if (Entries.Count == 0)
				{
					throw new ArgumentException();
				}

				Atom[] atoms = entries.SelectMany(e => e.Composition.Atoms).Distinct().Where(a => a != Atom.Vacancy).ToArray();

				if (atoms.Length != 2)
				{
					throw new ArgumentException("Only 2-atom systems are supported.");
				}

				// The non-vacancy atom and the corresponding per-unit cell counts.
				atom1 = atoms[0];
				atom2 = atoms[1];
				counts1 = entries.Select(p => p.Composition.Count(atom1)).ToArray();
				counts2 = entries.Select(p => p.Composition.Count(atom2)).ToArray();
				count1_min = counts1.Min();
				count1_max = counts1.Max();
				count2_min = counts2.Min();
				count2_max = counts2.Max();
		
				// Rolling cache of solvers different temperatures and pressures.
				solvers = new RollingCache<HashedTuple<double, double>, InnerSolver>(1000);
			}

			public double FreeEnergy(double temperature, double pressure, Composition composition)
			{
				double x1, x2;
				if (!Count(composition, out x1, out x2))
				{
					return double.NaN;
				}

				return CreateSolver(temperature, pressure).GibbsPotential(x1, x2);
			}

			public double Volume(double temperature, double pressure, Composition composition)
			{
				throw new NotImplementedException();
			}

			public ChemicalPotential ChemicalPotential(double temperature, double pressure, Composition composition)
			{
				double x1, x2;
				if (!Count(composition, out x1, out x2))
				{
					return new ChemicalPotential(atom1, double.NaN, atom2, double.NaN);
				}

				double mu1, mu2;
				CreateSolver(temperature, pressure).ChemicalPotential(x1, x2, out mu1, out mu2);

				return new ChemicalPotential(atom1, mu1, atom2, mu2);
			}

			public double VolumetricThermalExpansionCoefficient(double temperature, double pressure, Composition composition)
			{
				throw new NotImplementedException();
			}
			
			private bool Count(Composition composition, out double count1, out double count2)
			{
				if (composition == null)
				{
					throw new ArgumentNullException();
				}

				foreach (Atom a in composition.Atoms)
				{
					if (a == Atom.Vacancy || a == Atom.Undefined)
					{
						// Ignore vacancies.
						continue;
					}

					if (a != atom1 && a != atom2)
					{
						throw new ArgumentException();
					}
				}

				count1 = composition.Count(atom1);
				count2 = composition.Count(atom2);

				// Must be inside the convex hull spanned by the composition of the entries. Assuming that's just a triangle here.
				return count1 >= count1_min && count1 <= count1_max && count2 >= count2_min && count1 + count2 <= Math.Max(count1_max, count2_max);
			}

			private InnerSolver CreateSolver(double temperature, double pressure)
			{
				if (temperature < 0.0)
				{
					throw new ArgumentOutOfRangeException();
				}

				InnerSolver solver;
				if (!solvers.TryGetValue(HashedTuple.Create<double, double>(temperature, pressure), out solver))
				{
					solver = new InnerSolver(this, temperature, pressure);
					solvers.Add(HashedTuple.Create<double, double>(temperature, pressure), solver);
				}

				return solver;
			}

			public ReadOnlyList<IIsobaricPartitionFunctionEntry> Entries
			{
				get;
				private set;
			}

			IHelmholtzGrandPotential IGibbsGrandPotential.HelmholtzPotential
			{
				get
				{
					// Not really needed.
					throw new NotImplementedException();
				}
			}

			[Serializable]
			private class InnerSolver
			{
				private BivariateIsobaricPartitionFunction parent;
				private int n;
				private double[] energies, volumes;
				private double beta;
				private RollingCache<HashedTuple<double, double>, Tuple<double, double>> mus;

				public InnerSolver(BivariateIsobaricPartitionFunction parent, double t, double p)
				{
					this.parent = parent;

					n = parent.Entries.Count;

					// Precompute free energies for the given fixed temperature and pressure.
					energies = new double[n];
					volumes = new double[n];
					for (int i = 0; i < n; i++)
					{
						energies[i] = parent.Entries[i].Potential.FreeEnergy(t, p);

						// Postpone computation of volumes until requested?
						volumes[i] = parent.Entries[i].Potential.Volume(t, p);
					}

					double kb = 8.617332478e-5; // eV/K
					beta = 1.0 / (kb * t);

					// Rolling cache of chemical potentials.
					mus = new RollingCache<HashedTuple<double, double>, Tuple<double, double>>(10000);
				}

				private double LogPartitionFunction(double mu1, double mu2, out double[] f)
				{
					double[] e = new double[n];
					double e0 = double.NegativeInfinity;
					for (int i = 0; i < n; i++)
					{
						e[i] = beta * mu1 * parent.counts1[i] + beta * mu2 * parent.counts2[i] - beta * energies[i];
						e0 = Math.Max(e0, e[i]);
					}

					// Compute partition function with prefactor introduced for numerical stability.
					double z = 0.0;
					f = new double[n];
					for (int i = 0; i < n; i++)
					{
						f[i] = Math.Exp(e[i] - e0) * parent.Entries[i].Degeneracy;
						z += f[i];
					}

					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						f[i] /= z;
					}

					// Divide out prefactor again after taking the logarithm.
					double lz = Math.Log(z) + e0;

					return lz;
				}

				private DualNumber LogPartitionFunction(DualNumber mu1, DualNumber mu2, out DualNumber[] f)
				{
					DualNumber[] e = new DualNumber[n];
					DualNumber e0 = double.NegativeInfinity;
					for (int i = 0; i < n; i++)
					{
						e[i] = beta * mu1 * parent.counts1[i] + beta * mu2 * parent.counts2[i] - beta * energies[i];
						e0 = Max(e0, e[i]);
					}

					// Compute partition function with prefactor introduced for numerical stability.
					DualNumber z = 0.0;
					f = new DualNumber[n];
					for (int i = 0; i < n; i++)
					{
						f[i] = DualNumber.Exp(e[i] - e0) * parent.Entries[i].Degeneracy;
						z += f[i];
					}

					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						f[i] /= z;
					}

					// Divide out prefactor again after taking the logarithm.
					DualNumber lz = DualNumber.Log(z) + e0;

					return lz;
				}
				
				public double GrandPotential(double mu1, double mu2)
				{
					double[] f;
					double lz = LogPartitionFunction(mu1, mu2, out f);

					double phi = -lz / beta;

					return phi;
				}

				public void Count(double mu1, double mu2, out double x1, out double x2)
				{
					double[] f;
					LogPartitionFunction(mu1, mu2, out f);

					x1 = 0.0;
					x2 = 0.0;
					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						x1 += f[i] * parent.counts1[i];
						x2 += f[i] * parent.counts2[i];
					}
				}

				public void Count(DualNumber mu1, DualNumber mu2, out DualNumber x1, out DualNumber x2)
				{
					DualNumber[] f;
					LogPartitionFunction(mu1, mu2, out f);

					x1 = 0.0;
					x2 = 0.0;
					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						x1 += f[i] * parent.counts1[i];
						x2 += f[i] * parent.counts2[i];
					}
				}

				public void ChemicalPotential(double x1, double x2, out double mu1, out double mu2)
				{
					SolveChemicalPotential(x1, x2, out mu1, out mu2);
				}

				private void SolveChemicalPotential(double x1, double x2, out double mu1, out double mu2)
				{
					//SolveChemicalPotentialOptimizer(x1, x2, out mu1, out mu2);

					Tuple<double, double> mu;
					if (mus.TryGetValue(HashedTuple.Create<double, double>(x1, x2), out mu))
					{
						mu1 = mu.Item1;
						mu2 = mu.Item2;
					}
					else
					{
						SolveChemicalPotentialOptimizer(x1, x2, out mu1, out mu2);
						mus.Add(HashedTuple.Create<double, double>(x1, x2), Tuple.Create<double, double>(mu1, mu2));
					}
				}

				private void SolveChemicalPotentialOptimizer(double x1_0, double x2_0, out double mu1_0, out double mu2_0)
				{
					Variable mu1 = new Variable();
					Variable mu2 = new Variable();

					Function x1 = DualNumberFunction.Create(t => { DualNumber x1_, x2_; Count(t[mu1], t[mu2], out x1_, out x2_); return x1_; }, mu1, mu2);
					Function x2 = DualNumberFunction.Create(t => { DualNumber x1_, x2_; Count(t[mu1], t[mu2], out x1_, out x2_); return x2_; }, mu1, mu2);

					// Find the simultaneous zeroes as the as the minimum of this convex function. (Use the subgradient method?)
					Function f = Function.Sqr(x2 - x2_0) + Function.Sqr(x1 - x1_0);

					IpoptOptimizer o = new IpoptOptimizer();
					o.Variables.Add(mu1);
					o.Variables.Add(mu2);
					o.ObjectiveFunction = f;
					o.PrintLevel = 0;

					// Try a number of different initial chemical potentials. Reproducible as we're using the same seed!
					PolarBoxMullerGenerator r = new PolarBoxMullerGenerator(new LinearCongruentialGenerator());
					for (int i = 0; i < 20; i++)
					{
						// Usually the chemical potential per atom is around these values. Standard deviation is 1 here.
						IpoptOptimizerResult or = o.RunIpopt(mu1 | r.NextSample() - 8.0, mu2 | r.NextSample() - 8.0);

						if (!or.HasConverged)
						{
							continue;
						}

						mu1_0 = mu1.Value(or.OptimalPoint);
						mu2_0 = mu2.Value(or.OptimalPoint);

						// Verify result.
						double x1_2, x2_2;
						Count(mu1_0, mu2_0, out x1_2, out x2_2);

						if (Math.Abs(x1_0 - x1_2) > 1.0e-5 || Math.Abs(x2_0 - x2_2) > 1.0e-5)
						{
							continue;
						}

						// Optimal solution found.
						return;
					}

					// Exception here? Probably not!
					throw new Exception(string.Format("Couldn't find solution: x1 = {1}, x2 = {2}", x1_0, x2_0));

					// Couldn't find solution.
					//mu1_0 = double.NaN;
					//mu2_0 = double.NaN;
				}

				public double GibbsPotential(double x1, double x2)
				{
					// Legendre transform of the grand-potential.
					double mu1, mu2;
					ChemicalPotential(x1, x2, out mu1, out mu2);

					double phi = GrandPotential(mu1, mu2);

					double g = phi + mu1 * x1 + mu2 * x2;

					return g;
				}

				public double Volume(double x1, double x2)
				{
					double mu1, mu2;
					ChemicalPotential(x1, x2, out mu1, out mu2);

					double[] f;
					LogPartitionFunction(mu1, mu2, out f);

					double v = 0.0;
					for (int i = 0; i < n; i++)
					{
						// Prefactor doesn't matter for the Boltzmann factor.
						v += f[i] * volumes[i];
					}

					return v;
				}

				private static DualNumber Max(DualNumber f, DualNumber g)
				{
					return f.Value > g.Value ? f : g;
				}
			}
		}
	}
}
